import akshare as ak
import pandas as pd
import numpy as np
from config import ASSET_CODES, START_DATE, END_DATE, WINDOW_SIZE


def get_single_asset_data(asset_code, asset_type):
    """获取单资产数据（区分个股和ETF，适配不同接口参数）"""
    # 日期格式转换（用于筛选）
    start_date = pd.to_datetime(START_DATE)
    end_date = pd.to_datetime(END_DATE)

    print(f"获取{asset_type}数据：代码={asset_code}，起始日={START_DATE}，结束日={END_DATE}")

    try:
        if asset_type in ["股票ETF", "债券ETF"]:
            # ETF使用fund_etf_hist_sina接口（该接口不支持start/end_date，需后筛选）
            # 接口参数仅需symbol（含sh/sz前缀）
            df = ak.fund_etf_hist_sina(symbol=asset_code)

            # 打印ETF原始列名（确认列名）
            print(f"{asset_type}原始列名：{df.columns.tolist()}")

            # 将英文列名转换为中文列名（适配后续逻辑）
            df = df.rename(columns={
                'date': '日期',    # 接口返回英文date，转换为中文日期
                'open': '开盘',
                'high': '最高',
                'low': '最低',
                'close': '收盘',  # 接口返回close，转换为中文收盘
                'volume': '成交量' # 接口返回volume，转换为中文成交量
            })

            # 转换日期列并筛选时间范围（此时列名已改为“日期”，可正常操作）
            df["日期"] = pd.to_datetime(df["日期"])
            df = df[(df["日期"] >= start_date) & (df["日期"] <= end_date)]

        else:
            # 个股使用stock_zh_a_hist接口（支持start/end_date，需移除sh/sz前缀）
            asset_code_clean = asset_code.replace("sh", "").replace("sz", "")
            df = ak.stock_zh_a_hist(
                symbol=asset_code_clean,
                period="daily",
                start_date=START_DATE.replace("-", ""),
                end_date=END_DATE.replace("-", ""),
                adjust="qfq"
            )
            # 打印个股原始列名
            print(f"{asset_type}原始列名：{df.columns.tolist()}")

    except Exception as e:
        raise ConnectionError(f"【接口调用失败】{asset_type}({asset_code})：{str(e)}")

    # 校验数据是否为空
    if df.empty:
        raise ValueError(
            f"【数据为空错误】{asset_type}({asset_code}) 未获取到数据，可能原因：\n"
        )

    # 统一校验必要列
    required_cols = ["日期", "收盘", "成交量"]
    missing_cols = [col for col in required_cols if col not in df.columns]
    if missing_cols:
        raise KeyError(
            f"【列名错误】{asset_type}({asset_code}) 缺失必要列：{missing_cols}，\n"
            f"接口返回列名：{df.columns.tolist()}"
        )

    # 提取并转换列名
    df = df[required_cols].rename(columns={"收盘": "close", "成交量": "volume"})

    # 日期标准化并设置索引
    df["日期"] = pd.to_datetime(df["日期"])
    df.set_index("日期", inplace=True)
    return df


def preprocess_data(df):
    """数据预处理：计算MA/RSI、标准化、缺失值清洗"""
    # 打印预处理前的列名
    print(f"预处理前数据列名：{df.columns.tolist()}")

    df["ma5"] = df["close"].rolling(window=5).mean()
    df["ma20"] = df["close"].rolling(window=20).mean()
    df["rsi"] = calculate_rsi(df["close"], window=14)
    df["volume_pct"] = df["volume"].pct_change()
    df.dropna(inplace=True)

    # 检查特征列是否存在
    features = ["volume_pct", "ma5", "ma20", "rsi"]
    missing_features = [f for f in features if f not in df.columns]
    if missing_features:
        raise KeyError(f"【预处理错误】缺失特征列：{missing_features}")

    df[features] = (df[features] - df[features].mean()) / df[features].std()
    return df


def calculate_rsi(series, window=14):
    """计算RSI指标"""
    delta = series.diff(1)
    gain = (delta.where(delta > 0, 0)).rolling(window=window).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(window=window).mean()
    loss = loss.replace(0, 1e-8)  # 避免除以零
    rs = gain / loss
    rsi = 100 - (100 / (1 + rs))
    return rsi


def get_multi_asset_data():
    """多资产数据整合"""
    asset_dfs = []
    for asset_type, code in ASSET_CODES.items():
        print(f"\n正在处理{asset_type}（代码：{code}）...")
        df = get_single_asset_data(code, asset_type)
        df = preprocess_data(df)
        df.columns = [f"{col}_{asset_type}" for col in df.columns]
        asset_dfs.append(df)

    # 检查是否有数据
    if not asset_dfs:
        raise ValueError("【整合错误】未获取到任何资产数据")

    merged_df = asset_dfs[0]
    for df in asset_dfs[1:]:
        merged_df = pd.merge(merged_df, df, left_index=True, right_index=True, how="inner")

    # 检查合并后的数据
    if merged_df.empty:
        raise ValueError("【整合错误】多资产数据合并后为空，请检查日期范围是否有重叠")

    print(
        f"\n多资产数据共{len(merged_df)}个交易日，{len(merged_df.columns)}个特征")
    return merged_df


if __name__ == "__main__":
    try:
        data = get_multi_asset_data()
        print("\n最终数据预览：")
        print(data.head())
    except Exception as e:
        print(f"【运行错误】{str(e)}")